using System;
using System.Collections.Generic;

namespace Gj.Cache
{
    public class LFUCache<K, V> : ICache<K, V>
    {
        private int count;
        private readonly int capacity;
        private int minFrequency;
        private readonly Dictionary<int, LFUDoubleLinkedList<K, V>> frequencyMap;
        private readonly Dictionary<K, LFUNode<K, V>> map;
        private readonly Action<K, V> onRemove;

        public LFUCache(int capacity, Action<K, V> onRemoveCallback = null)
        {
            this.capacity = capacity;
            count = 0;
            map = new Dictionary<K, LFUNode<K, V>>(capacity);
            frequencyMap = new Dictionary<int, LFUDoubleLinkedList<K, V>>();
            minFrequency = int.MaxValue;

            onRemove = onRemoveCallback;
        }

        // get the value and then promote the node
        public bool Get(K key, out V value)
        {
            if (!map.ContainsKey(key))
            {
                value = default;
                return false;
            }

            var node = map[key];
            PromoteNode(node);
            value = node.Value;
            return true;
        }

        public bool Contains(K key)
        {
            return map.ContainsKey(key);
        }

        public void Add(K key, V value)
        {
            if (capacity == 0) return;

            // update the value of that node and then promote the node
            if (map.ContainsKey(key))
            {
                var node = map[key];
                node.Value = value;
                PromoteNode(node);
            }
            else
            {
                // remove the LFU node
                if (count == capacity)
                {
                    var dll = frequencyMap[minFrequency];
                    var lfuNode = dll.RemoveLFUNode();
                    if (dll.Count == 0) frequencyMap.Remove(minFrequency);
                    map.Remove(lfuNode.Key);
                    count--;
                    onRemove?.Invoke(lfuNode.Key, lfuNode.Value);
                }

                // add a new node
                var node = new LFUNode<K, V>(key, value);
                minFrequency = node.Frequency;
                map[key] = node;
                if (!frequencyMap.ContainsKey(node.Frequency))
                    frequencyMap[node.Frequency] = new LFUDoubleLinkedList<K, V>();
                frequencyMap[node.Frequency].AddToTop(node);
                count++;
            }
        }

        // remove target node from current frequency double-linked list,
        // and then promote it to higher frequency double-linked list
        private void PromoteNode(LFUNode<K, V> node)
        {
            var dll = frequencyMap[node.Frequency];
            dll.RemoveNode(node);
            if (dll.Count == 0)
            {
                frequencyMap.Remove(node.Frequency);
                if (minFrequency == node.Frequency) minFrequency++;
            }

            node.Frequency++;
            if (!frequencyMap.ContainsKey(node.Frequency))
                frequencyMap[node.Frequency] = new LFUDoubleLinkedList<K, V>();
            frequencyMap[node.Frequency].AddToTop(node);
        }
    }
}